import os
import json
from astropy.io import fits
from datetime import datetime, timedelta

# 数据目录
FITS_ROOT = "/data/nfsroot-vxpp/L1X/Imaging"
PREVIEW_ROOT = "/data/nfsroot-vxpp/L1X/preview/Imaging"
JSON_OUTPUT_DIR = "/data/vxpp_info/L1X/imgInfo"

# 判断图像类型
def classify_image(header):
    obs_type = header.get("OBS_TYPE", "").strip()
    workmode = header.get("WORKMODE", "").strip()
    naxis1 = header.get("NAXIS1", 0)
    naxis2 = header.get("NAXIS2", 0)

    if obs_type == "ToO-NOM-ACAL":
        if workmode == "dark":
            return "dark"
        elif workmode == "flat":
            return "flat"
        else:
            return "Imaging"
    else:
        return "fullimg" if naxis1 > 2000 and naxis2 > 2000 else "windowimg"

# 读取 FITS 头信息
def read_fits_metadata(fits_path):
    try:
        with fits.open(fits_path) as hdul:
            header = hdul[0].header
            return {
                "BGMEDIAN": header.get("BGMEDIAN", None),
                "BGSTD": header.get("BGSTD", None),
                "EE70": header.get("EE70", None),
                "EE70_RMS": header.get("EE70_RMS", None),
                "DATE-OBS": format_date_obs(header.get("DATE-OBS", "")),
                "EXPTIME": header.get("EXPTIME", None),
                "BAND": header.get("BAND", ""),
                "OBS_TYPE": header.get("OBS_TYPE", ""),
                "WORKMODE": header.get("WORKMODE", ""),
                "NAXIS1": header.get("NAXIS1", 0),
                "NAXIS2": header.get("NAXIS2", 0),
                "SUN_ANG": header.get("SUN_ANG", None),
                "MOON_ANG": header.get("MOON_ANG", None)
            }
    except Exception as e:
        print(f"Error reading {fits_path}: {e}")
        return None

# 格式化 DATE-OBS（去掉小数秒部分）
def format_date_obs(date_obs):
    if isinstance(date_obs, str) and "." in date_obs:
        return date_obs.split(".")[0]  # 去掉秒后的小数点部分
    return date_obs

# 处理每天的 FITS 图像
def process_fits_images(date_str):
    data = {t: {"type": t, "count": 0, "date": date_str, "images": []} for t in ["dark", "Imaging", "flat", "fullimg", "windowimg"]}

    date_dir = os.path.join(FITS_ROOT, date_str)
    if not os.path.exists(date_dir):
        print(f"No data found for {date_str}")
        return

    for target_name in os.listdir(date_dir):
        target_dir = os.path.join(date_dir, target_name)
        if not os.path.isdir(target_dir):
            continue

        for ii, fits_file in enumerate(os.listdir(target_dir)):
            if not fits_file.endswith(".fit"):
                continue
            
            #跳过B通道图像，只处理R通道
            if fits_file.find('_B_')>0:
                continue
            
            fits_path = os.path.join(target_dir, fits_file)
            metadata_R = read_fits_metadata(fits_path)
            fits_file_B = fits_file.replace('_R_', '_B_')
            fits_path_B = os.path.join(target_dir, fits_file_B)
            metadata_B = read_fits_metadata(fits_path_B)
            if not metadata_R or not metadata_B:
                continue
            
            img_type = classify_image(metadata_R)
            relative_fits_path = os.path.relpath(target_dir, FITS_ROOT)

            image_info = {
                "target": target_name,
                "fits_path": relative_fits_path,
                "image_name_b": fits_file.replace('_R_', '_B_'),
                "image_name_r": fits_file,
                "image_info_b": metadata_B,
                "image_info_r": metadata_R,
                "DATE-OBS": metadata_B.get("DATE-OBS", ""),
                "NAXIS1": metadata_B.get("NAXIS1", 0),
                "NAXIS2": metadata_B.get("NAXIS2", 0),
                "SUN_ANG": metadata_B.get("SUN_ANG", None),
                "MOON_ANG": metadata_B.get("MOON_ANG", None)
            }
            data[img_type]["images"].append(image_info)
            data[img_type]["count"] += 1
            
            # if ii > 20:
            #     break
    # **按 DATE-OBS 递增排序**
    for img_type in data:
        data[img_type]["images"].sort(key=lambda x: x.get("DATE-OBS", ""))
        
    return data

# 生成 JSON 文件
def save_json(data, date_str):
    year = date_str[:4]
    output_dir = os.path.join(JSON_OUTPUT_DIR, year)
    os.makedirs(output_dir, exist_ok=True)
    json_path = os.path.join(output_dir, f"{date_str}.json")

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
        
    return json_path

def processToday():
    json_path = ""
    today_str = datetime.today().strftime("%Y-%m-%d")
    fits_data = process_fits_images(today_str)
    if fits_data:
        json_path = save_json(fits_data, today_str)
        print(f"JSON saved: {today_str}.json")
        
    return json_path, today_str

def processYesterday(deltaDay=1):
    json_path = ""
    yesterday_str = (datetime.utcnow() - timedelta(days=deltaDay)).strftime("%Y-%m-%d")
    fits_data = process_fits_images(yesterday_str)
    if fits_data:
        json_path = save_json(fits_data, yesterday_str)
        print(f"JSON saved: {yesterday_str}.json")
        
    return json_path, yesterday_str

# 处理所有历史数据
def process_all_dates():
    if not os.path.exists(FITS_ROOT):
        print(f"Error: FITS root directory {FITS_ROOT} does not exist.")
        return
    
    # 获取所有日期文件夹，并按日期从新到旧排序
    date_dirs = sorted(
        [d for d in os.listdir(FITS_ROOT) if d[:4].isdigit() and len(d) == 10],  # 过滤出日期格式 YYYY-MM-DD
        reverse=True  # 从晚到早排序
    )

    today_str = datetime.today().strftime("%Y-%m-%d")
    for date_str in date_dirs:
        if date_str == today_str:
            print(f"Skipping {date_str}, it's the current day.")
            continue
        if date_str<="2025-03-20":
            continue
        json_path = os.path.join(JSON_OUTPUT_DIR, date_str[:4], f"{date_str}.json")
        if os.path.exists(json_path):
            print(f"Skipping {date_str}, JSON already exists.")
            continue  # 避免重复处理已生成的 JSON 文件
        
        print(f"Processing {date_str}...")
        fits_data = process_fits_images(date_str)
        if fits_data:
            save_json(fits_data, date_str)
            print(f"Saved JSON for {date_str}")
                
        # break

# 主函数
if __name__ == "__main__":
    process_all_dates()
